import sys
import os
import ollama
import json
import fitz  # PyMuPDF
from typing import List, Dict, Union, Any, Optional, Callable

try:
    # 尝试从上层目录的 base_code 导入
    parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    sys.path.append(parent_dir)
    from base_code import *
    from base_code.common import *
except ImportError:
    print("❌ 错误: 无法找到 'base_code' 模块。请确保 pocketflow 代码位于正确的路径。")

from local_tools import *  # Global Tools
from context_engineering import ContextEngine, ReasoningStep


class PromptManager:
    """
    集中管理所有提示词。
    """

    def get_perception_prompt(self) -> str:
        """生成用于理解用户初始意图的提示词。"""
        return """
        你的任务是分析用户的输入，并提取其完整目标。
        将这个目标总结成一个清晰、可执行的列表。
        例如:
        - 用户输入: "帮我查查根目录下 paper.pdf 这篇论文的第二页写了啥"
        - 你的输出: "（1）查找根目录下的 'paper.pdf'，（2）读取并返回 'paper.pdf' 文件第二页的内容"
        - 用户输入: "100的平方根乘以5是多少？"
        - 你的输出: "计算100的平方根并乘以5"

        请直接输出总结的目标，不要包含任何多余的解释。
        """

    def get_deliberation_prompt(self, goal: str, next_goal: str, history: Union[List[tuple], List[str]], tools_description: str,
                                experience_hint: str = "") -> str:
        """
        生成 ReAct 循环中的“思考”提示词。du
        支持压缩历史和经验提示，防止上下文崩溃。
        """
        history_str = ""
        for h in history[-4:]:
            history_str += h[0] + "\n" + h[1] + "\n" + h[2] + "\n"

        if next_goal is None: next_goal = goal

        # 经验提示部分
        experience_section = f"""
        **经验提示:**
        {experience_hint if experience_hint else "暂无相关经验"}
        """ if experience_hint else ""

        return f"""
        假设你是一个能够解决复杂问题的MIT工程学教授。你的任务是接收用户的请求，并通过一个“思考->行动->观察”的循环来完成它。
        你必须严格遵守以下规则和工作流程。

        **你的最终目标是:** {goal}
        **你的当前目标是:** {next_goal}

        **可用工具:**
        你只能使用以下定义的工具。不要编造任何不存在的工具。
        {tools_description}
        
         **重要规则:**
        1.  **语言一致性**: 你的"思考"和"最终答案"必须使用与用户请求相同的语言。如果用户用中文提问，你就必须用中文回答。
        2.  **当前目录**: 所有文件操作和终端命令都相对于当前工作目录执行。当前目录是 `{os.getcwd()}`。

        **思考与行动历史:**
        {history_str if history_str else "无 (这是第一步)"}
        {experience_section}

        **你的任务:**
        根据最终目标和历史记录，决定下一步行动。
        1.  **思考 (Thought):** 分析当前情况，制定下一步计划。
        2.  **行动 (Action):** 选择一个工具来执行计划，或者如果你认为任务已完成，则选择 'finish' 工具。
        3. 如果在思考过程中遇到缺少合适工具的情况，可以尝试自己完成总结、翻译、提炼大纲这类任务，否则可以"thought"中输出"大哥，工具不足啊！"并返回finish，

        **输出格式:**
        你的输出必须是一个严格的JSON对象，格式如下：
        {{
            "thought": "这里是你的思考过程...",
            "action": {{
                "name": "<tool_name>",
                "args": {{
                    "arg1": "value1",
                    "arg2": "value2"
                }}
            }}
        }}

        **特殊工具:**
        - 如果你认为已经获得了足够的信息来回答用户的目标，或者是当前工具列表没有满足要求的工具，请使用 `finish` 工具，并在 `args` 中提供最终答案:
          `{{"name": "finish", "args": {{"answer": "这里是最终答案"}}}}`

        现在，请根据目标和历史，生成你的下一步“思考”和“行动”。        
        """

    def get_judgment_prompt(
            self,
            goal: str,
            observation: str,
            tools_description: str,
            experience_hint: str = ""
    ) -> str:
        """
        用于任务进度评估的提示词。
        由 JudgmentNode 调用，以判断任务是否完成、并建议下一步目标。
        """

        return f"""
        你是一个任务评估智能体，负责判断当前任务是否完成。
    
        **输入:**
        - 当前任务目标 (goal): {goal}
        - 最近的观察结果 (observation): {observation}         
    
        **经验提示:**
        {experience_hint if experience_hint else "暂无经验提示"}
    
        **你的任务:**
        1. 判断任务是否已经完成；
        2. 如果未完成，请提出下一步的合理目标（例如“读取PDF摘要”或“分析摘要相关性”）；
        3. 如果任务已完成，请总结最终答案。
    
        **输出格式 (严格JSON):**
        {{
            "done": true 或 false,
            "next_goal": "新的任务目标",
            "final_answer": "如果任务完成，提供最终结果或总结，否则为''"
        }}
    
        请不要输出任何解释性文本，只输出JSON结果。
        
        **提示:**
        - 如果用户目标包含“摘要”、“分析”、“读取”等字样，而 observation 只显示文件列表，则任务未完成。
        - 如果 observation 中包含错误提示，请根据 history 判断是否需要更换工具或修改目标。
        - next_goal 应该是能直接用现有工具完成的小目标。
        """

    def get_finalize_prompt(
            self,
            goal: str,
            observation: str
    ) -> str:
        """
        用于最后的结果生成。
        """
        return f"""
        你是一位MIT的工程学教授，现在你负责判断当前智能体的返回的输出是否与学生问题中的要求一致，如果不一致，你需要帮助处理，并给出最终结果。
        
        **智能体输出**
        {observation}
        
        **学生问题***
        {goal}
        
        注意，仅仅处理并输出结果，不要做其它解释
        """

    def _compress_history_for_prompt(self, history: List[tuple]) -> str:
        """压缩历史记录，避免提示词过长"""
        if not history:
            return ""

        # 如果历史记录很短，直接返回
        if len(history) <= 3:
            return "\n".join([
                f"Thought: {thought}\nAction: {json.dumps(action)}\nObservation: {obs}"
                for thought, action, obs in history
            ])

        # 压缩策略：保留关键信息
        compressed = []

        # 保留第一步
        thought, action, obs = history[0]
        compressed.append(f"开始: Thought: {thought}\nAction: {json.dumps(action)}\nObservation: {obs[:500]}...")

        # 保留失败步骤
        for i, (thought, action, obs) in enumerate(history[1:-1], 1):
            if "错误" in str(obs).lower() or "失败" in str(obs).lower():
                compressed.append(f"步骤{i}: Thought: {thought[:30]}...\nAction: {json.dumps(action)}\nObservation: {obs[:50]}...")

        # 保留最后一步
        thought, action, obs = history[-1]
        compressed.append(f"最新: Thought: {thought}\nAction: {json.dumps(action)}\nObservation: {obs}")

        return "\n".join(compressed)


# ==============================================================================
#  Agent记忆：包含了个体记忆&上下文的引用
# ==============================================================================
class AgentMemory:
    """
    存储 Agent 工作流中所有状态的数据类。
    增强版本：集成上下文工程功能
    """

    def __init__(self):
        self.raw_input: Optional[str] = None
        self.goal: Optional[str] = None
        self.next_goal: Optional[str] = None
        self.history: List[tuple] = []  # (thought, action, observation)

        # Current step's data
        self.current_thought: Optional[str] = None
        self.current_action: Optional[Dict] = None
        self.current_observation: Any = None

        self.final_answer: Optional[str] = None
        self.error_message: Optional[str] = None

        # Control flags and counters
        self.iteration_count: int = 0
        self.max_iteration: int = 10   # 防止无限循环

        # 上下文工程相关
        self.context_engine: ContextEngine = ContextEngine()
        self.reasoning_steps: List[ReasoningStep] = []  # 用于上下文工程

    def reset_workflow_data(self):
        """在每次运行前重置临时数据。"""
        self.raw_input = None
        self.goal = None
        self.history = []
        self.current_thought = None
        self.current_action = None
        self.current_observation = None
        self.final_answer = None
        self.error_message = None
        self.iteration_count = 0
        self.reasoning_steps = []  # 重置推理步骤

    def add_reasoning_step(self, thought: str, action: str, result: str, success: bool = True):
        """添加推理步骤用于上下文分析"""
        step = ReasoningStep(thought, action, result, success)
        self.reasoning_steps.append(step)

    def get_compressed_context(self) -> Dict[str, Any]:
        """获取压缩后的上下文信息"""
        return self.context_engine.get_compressed_context(
            self.goal or "",
            self.reasoning_steps
        )

    def process_task_completion(self, success: bool) -> Dict[str, Any]:
        """处理任务完成，提取经验教训"""
        return self.context_engine.process_task(self.reasoning_steps, success)

    def add_history(self, thought: str, action: dict, observation: any):
        """向历史记录中添加一轮 ReAct 循环。"""
        history = []
        history.append(f" {thought}")
        action_str = json.dumps(action, ensure_ascii=False)
        history.append(f" {action_str}")
        history.append(f" {str(observation)}")
        self.history.append(tuple(history))


# ==============================================================================
#  工作流节点定义 (Nodes)
# ==============================================================================
class PerceptionNode(AsyncNode):
    """感知节点：分析用户原始输入，确立核心目标 (Goal)。"""
    def __init__(self, prompt_manager: PromptManager, ollama_tool: Callable):
        super().__init__()
        self.prompt_manager = prompt_manager
        self.ollama_tool = ollama_tool

    async def prep_async(self,shared):
        return shared

    async def exec_async(self, shared: AgentMemory):
        print("\n🤔 [Perception] 分析用户意图中...")
        try:
            goal = await self.ollama_tool(
                system_prompt=self.prompt_manager.get_perception_prompt(),
                user_prompt=shared.raw_input,
                as_json=False
            )
            shared.goal = goal.strip()
            print(f"🎯 确立目标: {shared.goal}")
            return "success"
        except Exception as e:
            shared.error_message = f"感知阶段出错: {e}"
            return "error"

    async def post_async(self,shared,prep_res,exec_res):
        return exec_res


class DeliberationNode(AsyncNode):
    """决策节点 (Reason): ReAct 循环的核心，负责思考并决定下一步行动。"""

    def __init__(self, prompt_manager: PromptManager, ollama_tool: Callable, registry: CommandRegistry):
        super().__init__()
        self.prompt_manager = prompt_manager
        self.ollama_tool = ollama_tool
        self.registry = registry

    async def prep_async(self,shared):
        return shared

    async def exec_async(self, shared: AgentMemory):
        print("\n🤔 [Deliberation] 思考下一步行动...")
        if len(shared.history) >= shared.max_iteration:
            shared.error_message = "已达到最大思考步数，任务终止。"
            return "error"

        # 获取压缩后的上下文信息
        compressed_context = shared.get_compressed_context()

        # 构建增强的提示词
        tools_desc = self.registry.get_tools_description()

        # 使用压缩后的历史和经验提示
        enhanced_history = compressed_context["compressed_history"]
        experience_hint = compressed_context["hint"]

        # 构建增强的提示词
        system_prompt = self.prompt_manager.get_deliberation_prompt(
            shared.goal,
            shared.next_goal,
            shared.history,
            tools_desc,
            experience_hint    # 添加经验提示
        )

        #print(system_prompt)

        try:
            response = await self.ollama_tool(
                system_prompt=system_prompt,
                user_prompt="请生成下一步的思考和行动。",
                as_json=True
            )

            shared.current_thought = response.get("thought", "没有产生有效思考。")
            shared.current_action = response.get("action")

            print(f"  - 思考: {shared.current_thought}")

            if not shared.current_action or "name" not in shared.current_action:
                raise ValueError("LLM 返回的 action 格式不正确。")

            action_name = shared.current_action["name"]
            print(f"  - 计划行动: {action_name}")

            if action_name == "finish":
                shared.final_answer = shared.current_action.get("args", {}).get("answer", "任务完成，但未提供明确答案。")
                return "finish"

            return "execute_tool"

        except Exception as e:
            shared.error_message = f"决策阶段出错: {e}"
            return "error"

    async def post_async(self, shared, prep_res, exec_res):
        return exec_res


class ToolExecutionNode(AsyncNode):
    """行动节点 (Act): 执行由 DeliberationNode 决定的工具。"""

    def __init__(self, registry: CommandRegistry):
        super().__init__()
        self.registry = registry

    async def prep_async(self,shared):
        return shared

    async def exec_async(self, shared: AgentMemory):
        action = shared.current_action
        tool_name = action["name"]
        tool_args = action.get("args", {})

        print(f"\n🛠️ [Execution] 执行工具: {tool_name} (参数: {tool_args})")

        try:
            tool_func = self.registry.get_tool(tool_name)

            # 异步执行工具
            maybe_coro = tool_func(**tool_args)
            if asyncio.iscoroutine(maybe_coro):
                result = await maybe_coro
            else:
                result = maybe_coro

            shared.current_observation = str(result)[:500]
            print(f"  - 观察结果: \n{str(result)}")  # 打印结果
            return "success"
        except KeyError:
            shared.current_observation = f"错误: 工具 '{tool_name}' 未找到。"
            return "success"  # 返回 success 以便 Judgment 节点记录错误并继续循环
        except Exception as e:
            shared.current_observation = f"错误: 执行工具 '{tool_name}' 时发生异常: {e}"
            return "success"  # 同上

    async def post_async(self, shared, prep_res, exec_res):
        return exec_res


class JudgmentNode(AsyncNode):
    """判断节点：评估工具执行结果，并将其计入历史和上下文工程。"""

    def __init__(self, prompt_manager: PromptManager, ollama_tool: Callable, registry: CommandRegistry):
        super().__init__()
        self.prompt_manager = prompt_manager
        self.ollama_tool = ollama_tool
        self.registry = registry

    async def prep_async(self, shared):
        return shared

    async def exec_async(self, shared: AgentMemory):
        print("\n⚖️ [Judgment] 评估结果并更新历史...")

        # 记录到标准历史
        shared.add_history(shared.current_thought, shared.current_action, shared.current_observation)
        action_name = shared.current_action.get("name", "unknown")
        success = not ("错误" in str(shared.current_observation) or "失败" in str(shared.current_observation))
        shared.add_reasoning_step(shared.current_thought or "", action_name, str(shared.current_observation)[:200],
                                  success)

        # 构建 LLM 评估输入
        eval_prompt = self.prompt_manager.get_judgment_prompt(
            goal=shared.goal,
            observation=shared.current_observation,
            tools_description=self.registry.get_tools_description(),
            experience_hint=shared.get_compressed_context().get("hint", "")
        )

        # print(f"eval_prompt: {eval_prompt}")

        # 历史压缩（用于上下文判断）
        short_history = shared.context_engine.curator.get_recent_history_text() \
            if hasattr(shared.context_engine.curator, "get_recent_history_text") else ""

        user_prompt = shared.goal

        # 调用 LLM 判断任务进度
        try:
            llm_response = await self.ollama_tool(
                system_prompt=eval_prompt,
                user_prompt=user_prompt,
                as_json=True
            )
        except Exception as e:
            print(f"⚠️ JudgmentNode 调用 LLM 出错: {e}")
            return "continue"

        # print(llm_response)
        done = llm_response.get("done", False)
        next_goal = llm_response.get("next_goal", "").strip()
        final_answer = llm_response.get("final_answer", "").strip()

        # 评估结果都会进入上下文累积
        #shared.context_engine.curator.add_feedback(f"[LLM评估] done={done}, next_goal={next_goal}")

        # 4️⃣ 根据 LLM 评估更新状态
        if done:
            shared.final_answer = final_answer or "任务完成，但模型未提供明确总结。"
            print(f"🏁 LLM 判断任务已完成: {shared.final_answer}")
            return "finish"
        elif next_goal:
            shared.next_goal = next_goal
            print(f"🎯 LLM 更新下一步目标: {shared.next_goal}")
            return "continue"
        else:
            print("ℹ️ LLM 未提供明确下一步，默认继续。")
            return "continue"

    async def post_async(self, shared, prep_res, exec_res):
        #print(f"exec_res: {exec_res}")
        return exec_res


class FinalizeNode(AsyncNode):
    """终结节点：在流程结束时整理并设置最终答案，同时处理上下文工程。"""

    def __init__(self, prompt_manager: PromptManager, ollama_tool: Callable):
        super().__init__()
        self.prompt_manager = prompt_manager
        self.ollama_tool = ollama_tool

    async def prep_async(self,shared):
        return shared

    async def exec_async(self, shared: AgentMemory):
        print("\n✅ [Finalize] 任务流程结束。")

        if shared.error_message:
            shared.final_answer = f"❌ 任务失败: {shared.error_message}"
            success = False
        elif not shared.final_answer:
            shared.final_answer = "ℹ️ 任务流程已完成，但没有生成最终答案。"
            success = False
        else:
            success = True

        # 完成最后整理
        final_prompt = self.prompt_manager.get_finalize_prompt(
            goal=shared.goal,
            observation=shared.current_observation
        )

        print(f"eval_prompt: {final_prompt}")

        # 历史压缩（用于上下文判断）
        short_history = shared.context_engine.curator.get_recent_history_text() \
            if hasattr(shared.context_engine.curator, "get_recent_history_text") else ""

        # user_prompt = json.dumps({
        #     "goal": shared.goal,
        #     "observation": str(shared.current_observation)[:500],
        #     "history": short_history
        # }, ensure_ascii=False)

        user_prompt = "根据目标和当前观察生成最终结果"

        # 调用 LLM 判断任务进度
        try:
            llm_response = await self.ollama_tool(
                system_prompt=final_prompt,
                user_prompt=user_prompt,
                as_json=True
            )
            print(llm_response)
        except Exception as e:
            print(f"⚠️ FinalizeNode 调用 LLM 出错: {e}")
            return


# ==============================================================================
#  4. ReActAgent (基于工作流的实现)
# ==============================================================================
class ReActAgent(AsyncFlow):
    """
    一个基于 pocketflow 工作流的 ReAct Agent。
    其内部逻辑由一个显式的“思考-行动-判断”循环构成。
    """

    def __init__(self, tools_system, llm="qwen3:4b"):
        super().__init__()

        # 1. 初始化核心组件
        self.registry = tools_system
        self.prompts = PromptManager()
        self.memory = AgentMemory()
        self.llm_model = llm

        # 2. 定义一个通用的 ollama 调用工具 (内部使用)
        # 这个工具不通过 @tool 注册，因为它被节点直接持有和调用
        async def _internal_ollama_call(system_prompt: str, user_prompt: str, as_json: bool = True):
            content = None
            try:
                response = await ollama.AsyncClient().chat(
                    model=self.llm_model,
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt},
                    ],
                    format="json" if as_json else None
                )
                content = response["message"]["content"]
                if as_json:
                    # <--- 建议: 增强 JSON 解析的鲁棒性
                    return json.loads(content)
                return content
            except json.JSONDecodeError as e:
                raise ValueError(f"LLM 未能返回有效的 JSON 格式。内容: '{content}'") from e
            except Exception as e:
                # 重新抛出异常，以便上层捕获
                raise e

        # 3. 定义工作流节点
        perception = PerceptionNode(self.prompts, _internal_ollama_call)
        deliberation = DeliberationNode(self.prompts, _internal_ollama_call, self.registry)
        executor = ToolExecutionNode(self.registry)
        judgment = JudgmentNode(self.prompts, _internal_ollama_call, self.registry)
        finalizer = FinalizeNode(self.prompts, _internal_ollama_call)

        # 4. 编排工作流
        self.start(perception)
        perception - "success" >> deliberation
        perception - "error" >> finalizer  # 感知失败直接结束

        deliberation - "execute_tool" >> executor
        deliberation - "finish" >> finalizer
        deliberation - "error" >> finalizer

        executor - "success" >> judgment  # 执行后总是进行判断

        judgment - "continue" >> deliberation  # 判断后返回思考，形成 ReAct 循环！
        judgment - "finish" >> finalizer
        judgment - "error" >> finalizer

        print("💡 ReActAgent (Workflow-based) 已初始化。")

    async def process_command_async(self, line: str) -> str:
        """Agent 的异步入口方法，处理用户输入。"""
        self.memory.reset_workflow_data()
        self.memory.raw_input = line

        await self.run_async(self.memory)

        return self.memory.final_answer or "流程结束，但没有最终答案。"

    def show_memo(self, verbose: bool = False):
        """打印当前Agent的记忆、洞察与统计信息（增强版）"""

        print("\n==============================")
        print("🧭 ReActAgent 状态调试面板")
        print("==============================")

        # 0️⃣ 当前状态概览
        print("🎯 当前任务目标:", self.memory.goal or "（未设定）")
        print("🔁 当前迭代次数:", self.memory.iteration_count, "/", self.memory.max_iteration)
        if self.memory.error_message:
            print("⚠️  错误信息:", self.memory.error_message)

        # 当前回合内容
        print("\n🧠 当前思考与执行状态:")
        print("  🪞 Thought:", (self.memory.current_thought or "（无）")[:200])
        print("  🧩 Action:", json.dumps(self.memory.current_action, ensure_ascii=False, indent=2)
        if self.memory.current_action else "（无）")
        print("  👁️  Observation:",
              (str(self.memory.current_observation)[:500] + "...")
              if self.memory.current_observation else "（无观察结果）")

        # 1️⃣ 上下文统计
        print("\n📊 上下文引擎统计:")
        stats = self.memory.context_engine.get_stats()
        print(f"  已处理任务数: {stats['tasks_processed']}")
        print(f"  当前洞察数量: {stats['total_insights']}")
        print(f"  估算记忆压缩率: {stats['memory_saved']}")

        # 2️⃣ 洞察回顾
        print("\n🧠 已提取的洞察 (Insights):")
        insights = self.memory.context_engine.curator.insights
        if not insights:
            print("  （暂无洞察）")
        else:
            for i, (trigger, ins) in enumerate(insights.items(), 1):
                print(f"  {i:02d}. [{trigger}] {ins.lesson} "
                      f"(权重={ins.weight:.1f}, 频次={ins.count})")

        # 3️⃣ 详细历史记录
        print("\n🕒 历史记录 (ReAct 回合):")
        if not self.memory.history:
            print("  （暂无历史记录）")
        else:
            for i, (thought, action, obs) in enumerate(self.memory.history[-5:], 1):
                print(f"  第 {i} 步：")
                print(f"    🧠 Thought: {thought[:200]}")
                print(f"    🧩 Action: {json.dumps(action, ensure_ascii=False)}")
                print(f"    👁️  Observation: {str(obs)[:300]}{'...' if len(str(obs)) > 300 else ''}")
                print("  -----------------------------")

        # 4️⃣ （可选）显示上下文引擎的内部步骤记忆
        if verbose:
            print("\n🪶 步骤压缩记忆 (ContextEngine):")
            steps = getattr(self.memory.context_engine.curator, "step_memory", [])
            if not steps:
                print("  （暂无步骤记忆）")
            else:
                for line in steps[-10:]:
                    print("  ", line)

        print("==============================\n")


